//
// Copyright 2019 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS-IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
//
package sched

import (
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"github.com/google/schedviz/tracedata/trace"
)

// ThreadTransitionBuilder instances are generated by ThreadTransitionSetBuilder
type ThreadTransitionBuilder struct {
	threadTransition *threadTransition
	stringBank       *stringBank
	err              error
}

// newThreadTransitionBuilder returns a new, empty ThreadTransitionBuilder.
func newThreadTransitionBuilder(eventID int, timestamp trace.Timestamp, pid PID, sb *stringBank) *ThreadTransitionBuilder {
	return &ThreadTransitionBuilder{
		threadTransition: &threadTransition{
			EventID:                  eventID,
			Timestamp:                timestamp,
			PID:                      pid,
			PrevCommand:              UnknownCommand,
			NextCommand:              UnknownCommand,
			PrevPriority:             UnknownPriority,
			NextPriority:             UnknownPriority,
			PrevCPU:                  UnknownCPU,
			NextCPU:                  UnknownCPU,
			PrevState:                AnyState,
			NextState:                AnyState,
			StatePropagatesThrough:   false,
			onForwardsStateConflict:  Fail,
			onBackwardsStateConflict: Fail,
			onForwardsCPUConflict:    Fail,
			onBackwardsCPUConflict:   Fail,
		},
		stringBank: sb,
	}
}

// WithPrevCommand sets the current threadTransition's PrevCommand.
func (ttb *ThreadTransitionBuilder) WithPrevCommand(prevCommand string) *ThreadTransitionBuilder {
	ttb.threadTransition.PrevCommand = ttb.stringBank.stringIDByString(prevCommand)
	return ttb
}

// WithNextCommand sets the current threadTransition's PrevCommand.
func (ttb *ThreadTransitionBuilder) WithNextCommand(nextCommand string) *ThreadTransitionBuilder {
	ttb.threadTransition.NextCommand = ttb.stringBank.stringIDByString(nextCommand)
	return ttb
}

// WithPrevPriority sets the current threadTransition's PrevPriority.
func (ttb *ThreadTransitionBuilder) WithPrevPriority(prevPriority Priority) *ThreadTransitionBuilder {
	ttb.threadTransition.PrevPriority = prevPriority
	return ttb
}

// WithNextPriority sets the current threadTransition's PrevPriority.
func (ttb *ThreadTransitionBuilder) WithNextPriority(nextPriority Priority) *ThreadTransitionBuilder {
	ttb.threadTransition.NextPriority = nextPriority
	return ttb
}

// WithPrevCPU sets the current threadTransition's PrevCPU.
func (ttb *ThreadTransitionBuilder) WithPrevCPU(prevCPU CPUID) *ThreadTransitionBuilder {
	ttb.threadTransition.PrevCPU = prevCPU
	return ttb
}

// WithNextCPU sets the current threadTransition's PrevCPU.
func (ttb *ThreadTransitionBuilder) WithNextCPU(nextCPU CPUID) *ThreadTransitionBuilder {
	ttb.threadTransition.NextCPU = nextCPU
	return ttb
}

// WithCPUPropagatesThrough sets whether CPUs can propagate through the
// current threadTransition, during CPU inferefence.
func (ttb *ThreadTransitionBuilder) WithCPUPropagatesThrough(cpuPropagatesThrough bool) *ThreadTransitionBuilder {
	ttb.threadTransition.CPUPropagatesThrough = cpuPropagatesThrough
	return ttb
}

// WithPrevState sets the current threadTransition's PrevState.
func (ttb *ThreadTransitionBuilder) WithPrevState(prevState ThreadState) *ThreadTransitionBuilder {
	if prevState&UnknownState != 0 {
		ttb.err = status.Errorf(codes.InvalidArgument, "error handling event %d: WithPrevState cannot include UnknownState", ttb.threadTransition.EventID)
	}
	ttb.threadTransition.PrevState = prevState
	return ttb
}

// WithNextState sets the current threadTransition's PrevState.
func (ttb *ThreadTransitionBuilder) WithNextState(nextState ThreadState) *ThreadTransitionBuilder {
	if nextState&UnknownState != 0 {
		ttb.err = status.Errorf(codes.InvalidArgument, "error handling event %d: WithNextState cannot include UnknownState", ttb.threadTransition.EventID)
	}
	ttb.threadTransition.NextState = nextState
	return ttb
}

// WithStatePropagatesThrough sets whether states can propagate through the
// current threadTransition during state inferefence.  This should be true for
// transitions that do not affect a thread's state (such as migrations) and
// false for transitions that do affect a thread's state.
func (ttb *ThreadTransitionBuilder) WithStatePropagatesThrough(statePropagatesThrough bool) *ThreadTransitionBuilder {
	ttb.threadTransition.StatePropagatesThrough = statePropagatesThrough
	return ttb
}

// OnForwardsStateConflict sets the ConflictPolicy to be used when this
// threadTransition's NextState conflicts with the next threadTransition's
// PrevState.
func (ttb *ThreadTransitionBuilder) OnForwardsStateConflict(policy ConflictPolicy) *ThreadTransitionBuilder {
	ttb.threadTransition.onForwardsStateConflict = policy
	return ttb
}

// OnBackwardsStateConflict sets the ConflictPolicy to be used when this
// threadTransition's PrevState conflicts with the previous threadTransition's
// NextState.
func (ttb *ThreadTransitionBuilder) OnBackwardsStateConflict(policy ConflictPolicy) *ThreadTransitionBuilder {
	ttb.threadTransition.onBackwardsStateConflict = policy
	return ttb
}

// OnForwardsCPUConflict sets the ConflictPolicy to be used when this
// threadTransition's NextCPU conflicts with the next threadTransition's
// PrevCPU.
func (ttb *ThreadTransitionBuilder) OnForwardsCPUConflict(policy ConflictPolicy) *ThreadTransitionBuilder {
	ttb.threadTransition.onForwardsCPUConflict = policy
	return ttb
}

// OnBackwardsCPUConflict sets the ConflictPolicy to be used when this
// threadTransition's PrevCPU conflicts with the previous threadTransition's
// NextCPU.
func (ttb *ThreadTransitionBuilder) OnBackwardsCPUConflict(policy ConflictPolicy) *ThreadTransitionBuilder {
	ttb.threadTransition.onBackwardsCPUConflict = policy
	return ttb
}

// ThreadTransitionSetBuilder instances are provided to event loader functions to
// facilitate their assembling correct threadTransitions.
type ThreadTransitionSetBuilder struct {
	stringBank *stringBank
	builders   []*ThreadTransitionBuilder
}

// newThreadTransitionSetBuilder returns a new, empty ThreadTransitionSetBuilder.
func newThreadTransitionSetBuilder(sb *stringBank) *ThreadTransitionSetBuilder {
	return &ThreadTransitionSetBuilder{
		stringBank: sb,
	}
}

// WithTransition starts a new threadTransition in the receiver, populated with
// the provided event ID, timestamp and PID, and with default values: unknown
// for all state, and false for all drops.
func (ttsb *ThreadTransitionSetBuilder) WithTransition(eventID int, timestamp trace.Timestamp, pid PID) *ThreadTransitionBuilder {
	ttb := newThreadTransitionBuilder(eventID, timestamp, pid, ttsb.stringBank)
	ttsb.builders = append(ttsb.builders, ttb)
	return ttb
}

func (ttsb *ThreadTransitionSetBuilder) transitions() ([]*threadTransition, error) {
	var ret = []*threadTransition{}
	for _, ttb := range ttsb.builders {
		ret = append(ret, ttb.threadTransition)
		if ttb.err != nil {
			return nil, ttb.err
		}
	}
	return ret, nil
}
